// Copyright 2019 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <fbl/auto_call.h>
#include <fuchsia/hardware/block/c/fidl.h>
#include <lib/fzl/fdio.h>
#include <ramdevice-client/ramdisk.h>
#include <unittest/unittest.h>
#include <zircon/assert.h>

#include "libgpt-tests.h"

// generate a random number between [1, max]
uint64_t random_non_zero_length(uint64_t max) { return (rand() % max) + 1; }

extern bool gUseRamDisk;
extern char gDevPath[PATH_MAX];

namespace {

using gpt::GptDevice;

// TODO(auradkar): consolidate this guid_t definition with one in
// ulib/gpt/gpt.cpp.
struct guid_t {
  uint32_t data1;
  uint16_t data2;
  uint16_t data3;
  uint8_t data4[8];
};

constexpr guid_t kGuid = {0x0, 0x1, 0x2, {0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xa}};

constexpr uint64_t partition_size(const gpt_partition_t* p) { return p->last - p->first + 1; }

bool destroy_gpt(int fd, uint64_t block_size, uint64_t offset, uint64_t block_count) {
  BEGIN_HELPER;

  char zero[block_size];
  memset(zero, 0, sizeof(zero));

  ASSERT_GT(block_count, 0, "Block count should be greater than zero");
  ASSERT_GT(block_size, 0, "Block count should be greater than zero");

  uint64_t first = offset;
  uint64_t last = offset + block_count - 1;

  for (uint64_t i = first; i <= last; i++) {
    ASSERT_EQ(pwrite(fd, zero, sizeof(zero), block_size * i), (ssize_t)sizeof(zero),
              "Failed to pwrite");
  }
  // fsync is not supported in rpc-server.cpp
  // TODO(ZX-3294) to fix this
  // ASSERT_EQ(fsync(fd), 0, "Failed to fsync");
  END_HELPER;
}

class Partitions {
 public:
  Partitions(uint32_t count, uint64_t first, uint64_t last);

  // Returns partition at index index. Returns null if index is out of range.
  const gpt_partition_t* GetPartition(uint32_t index) const;

  // Returns number of parition created/removed.
  uint32_t GetCount() const { return partition_count_; }

  // Marks a partition as created in GPT.
  bool MarkCreated(uint32_t index);

  // Mark a partition as removed in GPT.
  bool ClearCreated(uint32_t index);

  // Returns true if the GPT should have the partition.
  bool IsCreated(uint32_t index) const;

  // Returns number of partition that should exist on GPT.
  uint32_t CreatedCount() const;

  // Returns true if two partitions are the same.
  bool Compare(const gpt_partition_t* in_mem_partition,
               const gpt_partition_t* on_disk_partition) const;

  // Returns true if the partition p exists in partitions_.
  bool Find(const gpt_partition_t* p, uint32_t* out_index) const;

 private:
  // List of partitions
  gpt_partition_t partitions_[gpt::kPartitionCount];

  // A variable to track whether a partition is created on GPT or not
  bool created_[gpt::kPartitionCount] = {};

  // Number of partitions_ that is populated with valid information
  uint32_t partition_count_;
};

Partitions::Partitions(uint32_t count, uint64_t first, uint64_t last) {
  ZX_ASSERT(count > 0);
  ZX_ASSERT(count <= gpt::kPartitionCount);
  partition_count_ = count;
  guid_t guid = kGuid;

  uint64_t part_first = first, part_last;
  uint64_t part_max_len = (last - first) / partition_count_;
  ZX_ASSERT(part_max_len > 0);

  memset(partitions_, 0, sizeof(partitions_));
  for (uint32_t i = 0; i < partition_count_; i++) {
    part_last = part_first + random_non_zero_length(part_max_len);
    guid.data1 = i;
    memcpy(partitions_[i].type, &guid, sizeof(partitions_[i].type));
    memcpy(partitions_[i].guid, &guid, sizeof(partitions_[i].type));
    partitions_[i].first = part_first;
    partitions_[i].last = part_last;
    partitions_[i].flags = 0;
    memset(partitions_[i].name, 0, sizeof(partitions_[i].name));
    snprintf(reinterpret_cast<char*>(partitions_[i].name), sizeof(partitions_[i].name), "%u_part",
             i);

    part_first += part_max_len;
  }
}

const gpt_partition_t* Partitions::GetPartition(uint32_t index) const {
  if (index >= partition_count_) {
    return nullptr;
  }

  return &partitions_[index];
}

bool Partitions::MarkCreated(uint32_t index) {
  BEGIN_HELPER;
  ASSERT_LT(index, partition_count_, "Index out of range");

  created_[index] = true;
  END_HELPER;
}

bool Partitions::ClearCreated(uint32_t index) {
  BEGIN_HELPER;
  ASSERT_LT(index, partition_count_, "Index out of range");

  created_[index] = false;
  END_HELPER;
}

bool Partitions::IsCreated(uint32_t index) const {
  BEGIN_HELPER;
  ASSERT_LT(index, partition_count_, "Index out of range");

  return created_[index];
  END_HELPER;
}

uint32_t Partitions::CreatedCount() const {
  uint32_t created_count = 0;

  for (uint32_t i = 0; i < GetCount(); i++) {
    if (IsCreated(i)) {
      created_count++;
    }
  }
  return (created_count);
}

bool Partitions::Compare(const gpt_partition_t* in_mem_partition,
                         const gpt_partition_t* on_disk_partition) const {
  if (memcmp(in_mem_partition->type, on_disk_partition->type, sizeof(in_mem_partition->type)) !=
      0) {
    return false;
  }

  if (memcmp(in_mem_partition->guid, on_disk_partition->guid, sizeof(in_mem_partition->guid)) !=
      0) {
    return false;
  }

  if (in_mem_partition->first != on_disk_partition->first) {
    return false;
  }

  if (in_mem_partition->last != on_disk_partition->last) {
    return false;
  }

  if (in_mem_partition->flags != on_disk_partition->flags) {
    return false;
  }

  // In mem partition name is a c-string whereas on-disk partition name
  // is stored as UTF-16. We need to convert UTF-16 to c-string before we
  // compare.
  char name[GPT_NAME_LEN];
  memset(name, 0, GPT_NAME_LEN);
  utf16_to_cstring(name, (const uint16_t*)on_disk_partition->name, GPT_NAME_LEN / 2);

  if (strncmp(name, reinterpret_cast<const char*>(in_mem_partition->name), GPT_NAME_LEN / 2) != 0) {
    return false;
  }

  return true;
}

bool Partitions::Find(const gpt_partition_t* p, uint32_t* out_index) const {
  for (uint32_t i = 0; i < partition_count_; i++) {
    if (Compare(GetPartition(i), p)) {
      *out_index = i;
      return true;
    }
  }

  return false;
}

// Defines a libgpt test function which can be passed to the TestWrapper.
typedef bool (*TestFunction)(LibGptTest* libGptTest);
#define RUN_TEST_WRAP(test_name) RUN_TEST_MEDIUM(TestWrapper<test_name>)

// A test wrapper which runs a libgpt test.
template <TestFunction TestFunc>
bool TestWrapper(void) {
  BEGIN_TEST;

  LibGptTest libGptTest(gUseRamDisk);
  ASSERT_TRUE(libGptTest.Init(), "Setting up the block device");
  // Run the test. This should pass.
  ASSERT_TRUE(TestFunc(&libGptTest));
  ASSERT_TRUE(libGptTest.Teardown(), "Tearing down and cleaning up the block device");

  END_TEST;
}

bool LibGptTest::Reset() {
  BEGIN_HELPER;
  fbl::unique_ptr<GptDevice> gpt;

  // explicitly close the fd, if open, before we attempt to reopen it.
  fd_.reset();

  fd_.reset(open(disk_path_, O_RDWR));

  ASSERT_TRUE(fd_.is_valid(), "Could not open block device\n");
  ASSERT_EQ(GptDevice::Create(fd_.get(), GetBlockSize(), GetBlockCount(), &gpt), ZX_OK);
  gpt_ = std::move(gpt);
  END_HELPER;
}

bool LibGptTest::Finalize() {
  BEGIN_HELPER;
  ASSERT_FALSE(gpt_->Valid(), "Valid GPT on uninitialized disk");

  ASSERT_EQ(gpt_->Finalize(), ZX_OK, "Failed to finalize");
  ASSERT_TRUE(gpt_->Valid(), "Invalid GPT after finalize");
  END_HELPER;
}

bool LibGptTest::Sync() {
  BEGIN_HELPER;

  ASSERT_EQ(gpt_->Sync(), ZX_OK, "Failed to sync");
  ASSERT_TRUE(gpt_->Valid(), "Invalid GPT after sync");

  END_HELPER;
}

bool LibGptTest::ReadRange() {
  BEGIN_HELPER;

  ASSERT_EQ(gpt_->Range(&usable_start_block_, &usable_last_block_), ZX_OK,
            "Retrieval of device range failed.");

  // TODO(auradkar): GptDevice doesn't export api to get GPT-metadata size.
  // If it does, we can keep better track of metadata size it says it needs
  // and metadata it actually uses.
  ASSERT_LT(GetUsableStartBlock(), GetBlockCount(), "Range starts after EOD");
  ASSERT_LT(GetUsableStartBlock(), GetUsableLastBlock(), "Invalid range");
  ASSERT_LT(GetUsableLastBlock(), GetBlockCount(), "Range end greater than block count");
  ASSERT_GT(GetUsableBlockCount(), 0, "GPT occupied all available blocks");

  END_HELPER;
}

bool LibGptTest::PrepDisk(bool sync) {
  BEGIN_HELPER;

  if (sync) {
    Sync();
  } else {
    Finalize();
  }

  ASSERT_TRUE(ReadRange(), "Read range failed");
  END_HELPER;
}

bool LibGptTest::InitDisk(const char* disk_path) {
  BEGIN_HELPER;

  use_ramdisk_ = false;
  snprintf(disk_path_, PATH_MAX, "%s", disk_path);
  fbl::unique_fd fd(open(disk_path_, O_RDWR));
  ASSERT_TRUE(fd.is_valid(), "Could not open block device to fetch info");
  fuchsia_hardware_block_BlockInfo block_info;
  fzl::UnownedFdioCaller caller(fd.get());
  zx_status_t status;
  ASSERT_EQ(fuchsia_hardware_block_BlockGetInfo(caller.borrow_channel(), &status, &block_info),
            ZX_OK);
  ASSERT_EQ(status, ZX_OK);

  blk_size_ = block_info.block_size;
  blk_count_ = block_info.block_count;

  ASSERT_GE(GetDiskSize(), kAccptableMinimumSize, "Insufficient disk space for tests");
  fd_ = std::move(fd);

  END_HELPER;
}

bool LibGptTest::InitRamDisk() {
  BEGIN_HELPER;
  ASSERT_EQ(ramdisk_create(GetBlockSize(), GetBlockCount(), &ramdisk_), ZX_OK,
            "Could not create ramdisk");
  strlcpy(disk_path_, ramdisk_get_path(ramdisk_), sizeof(disk_path_));
  fd_.reset(open(disk_path_, O_RDWR));
  if (!fd_) {
    return false;
  }

  END_HELPER;
}

bool LibGptTest::Init() {
  BEGIN_HELPER;
  auto error = fbl::MakeAutoCall([this]() { Teardown(); });
  if (use_ramdisk_) {
    ASSERT_TRUE(InitRamDisk());
  } else {
    ASSERT_TRUE(InitDisk(gDevPath));
  }

  // TODO(auradkar): All tests assume that the disks don't have an initialized
  // disk. If tests find an GPT initialized disk at the beginning of test,
  // they fail. The tests leave disks in initalized state.
  //
  // To either uninitialize an initialized disk as a part of setup
  // test needs to know where gpt lies on the disk. As of now libgpt doesn't
  // export an api to get the location(s) of the gpt on disk. So, we assume
  // here that gpt lies in first few (GptMetadataBlocksCount()) blocks on the
  // device. We also ignore any backup copies on the device.
  // Once there exists an api in libgpt to get size and location(s) of gpt,
  // we can setup/cleanup before/after running tests in a better way.
  ASSERT_TRUE(destroy_gpt(fd_.get(), GetBlockSize(), 0, GptMetadataBlocksCount()),
              "Failed to destroy gpt");

  ASSERT_TRUE(Reset());
  error.cancel();
  END_HELPER;
}

bool LibGptTest::TearDownDisk() {
  BEGIN_HELPER;
  ASSERT_FALSE(use_ramdisk_);
  END_HELPER;
}

bool LibGptTest::TearDownRamDisk() {
  BEGIN_HELPER;
  ASSERT_EQ(ramdisk_destroy(ramdisk_), ZX_OK);
  END_HELPER;
}

bool LibGptTest::Teardown() {
  BEGIN_HELPER;

  if (use_ramdisk_) {
    ASSERT_TRUE(TearDownRamDisk());
  } else {
    ASSERT_TRUE(TearDownDisk());
  }

  END_HELPER;
}

}  // namespace

// Creates "partitions->GetCount()"" number of partitions on GPT.
//  The information needed to create a partitions is passed in "partitions".
bool AddPartitionHelper(LibGptTest* libGptTest, Partitions* partitions) {
  BEGIN_HELPER;

  ASSERT_GT(partitions->GetCount(), 0, "At least one partition is required");
  for (uint32_t i = 0; i < partitions->GetCount(); i++) {
    const gpt_partition_t* p = partitions->GetPartition(i);
    ASSERT_EQ(libGptTest->AddPartition(reinterpret_cast<const char*>(p->name), p->type, p->guid,
                                       p->first, partition_size(p), p->flags),
              ZX_OK, "Add partition failed");
    partitions->MarkCreated(i);
  }

  END_HELPER;
}

// Removes randomly selected "remove_count" number of partitions.
bool RemovePartitionsHelper(LibGptTest* libGptTest, Partitions* partitions, uint32_t remove_count) {
  BEGIN_HELPER;
  uint32_t index;

  ASSERT_LE(remove_count, partitions->GetCount(), "Remove count exceeds whats available");
  ASSERT_LE(remove_count, partitions->CreatedCount(), "Cannot remove more partitions than created");

  for (uint32_t i = 0; i < remove_count; i++) {
    while (true) {
      index = static_cast<uint32_t>(rand()) % partitions->GetCount();
      if (partitions->IsCreated(index)) {
        break;
      }
    }
    ASSERT_TRUE(partitions->IsCreated(index), "Partition already removed");
    const gpt_partition_t* p = partitions->GetPartition(index);
    ASSERT_EQ(libGptTest->RemovePartition(p->guid), ZX_OK, "Failed to remove partition");
    partitions->ClearCreated(index);
  }
  END_HELPER;
}

// Verifies all the partitions that exists on GPT are the ones that are created
// by the test and vice-versa.
bool PartitionVerify(LibGptTest* libGptTest, const Partitions* partitions) {
  BEGIN_HELPER;
  bool found[gpt::kPartitionCount] = {};
  uint32_t found_index;

  // Check what's found on disk is created by us
  // iteratre over all partition that are present on disk and make sure
  // that we intended to create them.
  // Note: The index of an entry/partition need not match with the index of
  // the partition in "Partition* partition".
  for (uint32_t i = 0; i < gpt::kPartitionCount; i++) {
    gpt_partition_t* p = libGptTest->GetPartition(i);

    if (p == NULL) {
      continue;
    }

    ASSERT_TRUE(partitions->Find(p, &found_index), "Found an entry on GPT that we did not create");

    ASSERT_TRUE(partitions->IsCreated(found_index), "Removed entry reincarnated");
    found[found_index] = true;
  }

  // Check what's created is found on disk
  for (uint32_t i = 0; i < partitions->GetCount(); i++) {
    if (partitions->IsCreated(i)) {
      ASSERT_TRUE(found[i], "Created partition is missing on disk");
    }
  }

  END_TEST;
}

// Creates partitions and verifies them.
bool AddPartitions(LibGptTest* libGptTest, Partitions* partitions, bool sync) {
  BEGIN_HELPER;

  ASSERT_TRUE(AddPartitionHelper(libGptTest, partitions), "AddPartitionHelper failed");

  if (sync) {
    ASSERT_TRUE(libGptTest->Sync(), "Sync failed");
  }

  ASSERT_TRUE(PartitionVerify(libGptTest, partitions), "Partition verify failed");
  ASSERT_EQ(partitions->GetCount(), partitions->CreatedCount(),
            "Not as many created as we wanted to");

  END_HELPER;
}

// Removes partitions and verifies them.
bool RemovePartitions(LibGptTest* libGptTest, Partitions* partitions, uint32_t remove_count,
                      bool sync) {
  BEGIN_HELPER;

  ASSERT_TRUE(RemovePartitionsHelper(libGptTest, partitions, remove_count),
              "RemovePartitionsHelper failed");
  if (sync) {
    ASSERT_TRUE(libGptTest->Sync(), "Sync failed");
  }

  ASSERT_TRUE(PartitionVerify(libGptTest, partitions), "Partition verify failed");
  ASSERT_EQ(partitions->GetCount() - partitions->CreatedCount(), remove_count,
            "Not as many removed as we wanted to");

  END_HELPER;
}

// Removes all partitions and verifies them.
bool RemoveAllPartitions(LibGptTest* libGptTest, Partitions* partitions, bool sync) {
  BEGIN_HELPER;

  ASSERT_LE(partitions->GetCount(), partitions->CreatedCount(), "Not all partitions populated");
  ASSERT_EQ(libGptTest->RemoveAllPartitions(), ZX_OK, "Failed to remove all partition");

  for (uint32_t i = 0; i < partitions->GetCount(); i++) {
    partitions->ClearCreated(i);
  }

  ASSERT_TRUE(PartitionVerify(libGptTest, partitions), "Partition verify failed");
  ASSERT_EQ(partitions->CreatedCount(), 0, "Not as many removed as we wanted to");
  END_HELPER;
}

// Tests if we can create a GptDevice.
bool CreateTest(LibGptTest* libGptTest) {
  BEGIN_TEST;

  ASSERT_FALSE(libGptTest->IsGptValid(), "Valid GPT on uninitialized disk");
  ASSERT_TRUE(libGptTest->Reset(), "Failed to reset Test");
  ASSERT_FALSE(libGptTest->IsGptValid(), "Valid GPT after reset");
  END_TEST;
}

// Tests Finalize initializes GPT in-memory only and doesn't commit to disk.
bool FinalizeTest(LibGptTest* libGptTest) {
  BEGIN_TEST;

  ASSERT_TRUE(libGptTest->Finalize(), "Finalize failed");

  // Finalize initializes GPT but doesn't write changes to disk.
  // Resetting the Test should bring invalid gpt back.
  ASSERT_TRUE(libGptTest->Reset(), "Failed to reset Test");
  ASSERT_FALSE(libGptTest->IsGptValid(), "Valid GPT after finalize and reset");
  END_TEST;
}

// Tests Finalize initializes GPT and writes it to disk.
bool SyncTest(LibGptTest* libGptTest) {
  BEGIN_TEST;

  ASSERT_FALSE(libGptTest->IsGptValid(), "Valid GPT on uninitialized disk");

  // Sync should write changes to disk. Resetting should bring valid gpt back.
  ASSERT_TRUE(libGptTest->Sync(), "Sync failed");
  ASSERT_TRUE(libGptTest->Reset(), "Failed to reset Test");
  ASSERT_TRUE(libGptTest->IsGptValid(), "Invalid GPT after sync and reset");
  END_TEST;
}

// Tests the range the GPT blocks falls within disk.
bool RangeTest(LibGptTest* libGptTest) {
  BEGIN_TEST;

  ASSERT_TRUE(libGptTest->Finalize(), "Finalize failed");
  ASSERT_TRUE(libGptTest->ReadRange(), "Failed to read range");

  END_TEST;
}

// Test adding total_partitions partitions to GPT w/o writing to disk
template <uint32_t total_partitions, bool sync>
bool AddPartitionTest(LibGptTest* libGptTest) {
  BEGIN_TEST;
  ASSERT_TRUE(libGptTest->PrepDisk(sync), "Failed to setup disk");

  Partitions partitions(total_partitions, libGptTest->GetUsableStartBlock(),
                        libGptTest->GetUsableLastBlock());

  ASSERT_TRUE(AddPartitions(libGptTest, &partitions, sync), "AddPartitions failed");
  END_TEST;
}

// Test adding total_partitions partitions to GPT and test removing remove_count
// partitions later w/o writing to disk.
template <uint32_t total_partitions, uint32_t remove_count, bool sync>
bool RemovePartitionTest(LibGptTest* libGptTest) {
  BEGIN_TEST;
  ASSERT_TRUE(libGptTest->PrepDisk(sync), "Failed to setup disk");

  Partitions partitions(total_partitions, libGptTest->GetUsableStartBlock(),
                        libGptTest->GetUsableLastBlock());

  ASSERT_TRUE(AddPartitions(libGptTest, &partitions, sync), "AddPartitions failed");
  ASSERT_TRUE(RemovePartitions(libGptTest, &partitions, remove_count, sync),
              "RemovePartitions failed");

  END_TEST;
}

// Test removing all total_partititions from GPT w/o syncing.
template <uint32_t total_partitions, bool sync>
bool RemovePartitionAllTest(LibGptTest* libGptTest) {
  BEGIN_TEST;

  ASSERT_TRUE(libGptTest->PrepDisk(sync), "Failed to setup disk");

  Partitions partitions(total_partitions, libGptTest->GetUsableStartBlock(),
                        libGptTest->GetUsableLastBlock());

  ASSERT_TRUE(AddPartitions(libGptTest, &partitions, sync), "AddPartitions failed");
  ASSERT_TRUE(RemoveAllPartitions(libGptTest, &partitions, sync), "RemoveAllPartitions failed");

  END_TEST;
}

// Test if Diffs after adding partitions reflect all the changes.
template <uint32_t total_partitions>
bool DiffsTest(LibGptTest* libGptTest) {
  BEGIN_TEST;
  uint32_t diffs;

  ASSERT_NE(libGptTest->GetDiffs(0, &diffs), ZX_OK, "GetDiffs should fail before PrepDisk");
  ASSERT_TRUE(libGptTest->PrepDisk(false), "Failed to setup disk");
  ASSERT_NE(libGptTest->GetDiffs(0, &diffs), ZX_OK,
            "GetDiffs for non-existing partition should fail");

  Partitions partitions(total_partitions, libGptTest->GetUsableStartBlock(),
                        libGptTest->GetUsableLastBlock());
  ASSERT_TRUE(AddPartitions(libGptTest, &partitions, false), "AddPartitions failed");
  ASSERT_EQ(libGptTest->GetDiffs(0, &diffs), ZX_OK, "Diffs zero after adding partition");

  ASSERT_EQ(diffs,
            gpt::kGptDiffType | gpt::kGptDiffGuid | gpt::kGptDiffFirst | gpt::kGptDiffLast |
                gpt::kGptDiffName,
            "Unexpected diff after creating partition");
  ASSERT_TRUE(libGptTest->Sync(), "Failed to sync");
  ASSERT_EQ(libGptTest->GetDiffs(0, &diffs), ZX_OK, "GetDiffs failed");
  ASSERT_EQ(diffs, 0, "Diffs not zero after syncing partition");

  END_TEST;
}

BEGIN_TEST_CASE(libgpt_tests)
RUN_TEST_WRAP(CreateTest)
RUN_TEST_WRAP(FinalizeTest)
RUN_TEST_WRAP(SyncTest)
RUN_TEST_WRAP(RangeTest)
RUN_TEST_WRAP((AddPartitionTest<3, false>))
RUN_TEST_WRAP((AddPartitionTest<20, true>))
RUN_TEST_WRAP((RemovePartitionTest<12, 4, false>))
RUN_TEST_WRAP((RemovePartitionTest<3, 2, true>))
RUN_TEST_WRAP((RemovePartitionTest<11, 11, false>))
RUN_TEST_WRAP((RemovePartitionAllTest<12, true>))
RUN_TEST_WRAP((RemovePartitionAllTest<15, false>))
RUN_TEST_WRAP((DiffsTest<9>))
END_TEST_CASE(libgpt_tests);
